# vectors.py

import os
import logging
import requests
from langchain_community.document_loaders import UnstructuredMarkdownLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_huggingface import HuggingFaceEmbeddings
from pymilvus import MilvusClient
from langchain.schema import Document
from datetime import datetime

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# 支持的文件类型
SUPPORTED_EXTENSIONS = {
    '.pdf', '.doc', '.ppt', '.docx', '.pptx', '.html',
    '.txt', '.rtf', '.odt', '.xls', '.xlsx', '.md'
}

class EmbeddingsManager:
    def __init__(
        self,
        model_name: str = "moka-ai/m3e-base",
        device: str = "cpu",
        encode_kwargs: dict = {"normalize_embeddings": True},
        collection_name: str = "vector_db",
    ):
        """
        Initializes the EmbeddingsManager with the specified model and Milvus settings.

        Args:
            model_name (str): The HuggingFace model name for embeddings.
            device (str): The device to run the model on ('cpu' or 'cuda').
            encode_kwargs (dict): Additional keyword arguments for encoding.
            collection_name (str): The name of the Milvus collection.
        """
        self.model_name = model_name
        self.device = device
        self.encode_kwargs = encode_kwargs
        # 使用 Milvus Lite 模式（嵌入式）
        # 确保使用绝对路径
        self.milvus_db_file = "milvus_data/milvus.db"
        self.collection_name = collection_name
        
        self.embeddings = HuggingFaceEmbeddings(
            model_name=self.model_name,
            model_kwargs={"device": self.device},
            encode_kwargs=self.encode_kwargs,
        )
        
        # 确保数据库目录存在
        os.makedirs(os.path.dirname(self.milvus_db_file), exist_ok=True)

    # 处理文件
    def create_embeddings(self, file_path: str):
        """
        Processes the PDF, creates embeddings, and stores them in Milvus.

        Args:
            pdf_path (str): The file path to the PDF document.

        Returns:
            str: Success message upon completion.
        """
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"The file {file_path} does not exist.")

        # Load and preprocess the document
        # 使用 Unstructured 库加载文档，扩充到 PDF Word Excel PPT 
        file_ext = os.path.splitext(file_path)[-1].lower()
        if file_ext not in SUPPORTED_EXTENSIONS:
            raise ValueError(f"Unsupported file type: {file_ext}")
        
        # 如果不是 Makrdown 调用新文档转换服务器 API 进行文件转换
        if file_ext != '.md':
            # 调用新文档转换服务器 API 进行文件转换
            api_url = "http://localhost:5000/convert"
            payload = {
                "input_path": os.path.abspath(file_path),
                "method": "ocr",
            }
            response = requests.post(api_url, json=payload)
            result = response.json()
            print(f"Conversion result: {result}")

            if result.get('status') != 'success':
                raise Exception(f"Conversion failed: {result.get('message')}")
            
            # 获取转换后的 markdown 文件路径
            md_path = "./MinerU/db_files/{}.md".format(file_path.split('/')[-1].split('.')[0])
        else:
            md_path = file_path

        # 使用 Unstructured 库加载 markdown 文件
        loader = UnstructuredMarkdownLoader(md_path)
        docs = loader.load()
        if not docs:
            raise ValueError("No documents were loaded from the PDF.")

        # 提取文档标题和结构
        for doc in docs:
            # 从文件名获取标题
            filename = os.path.basename(file_path)
            title = os.path.splitext(filename)[0]
            
            # 从文档内容中提取可能的标题
            content = doc.page_content
            first_line = content.split('\n')[0].strip()
            if len(first_line) < 100 and not first_line.startswith(('#', '-', '*', '1.', '2.')):
                title = first_line
            
            # 更新元数据
            doc.metadata.update({
                'title': title,
                'filename': filename,
                'source': file_path,
                'document_type': file_ext[1:].upper(),
                'processing_time': datetime.now().isoformat()
            })

        # 使用更大的文本块进行分割
        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=400,  # 增大块大小
            chunk_overlap=200,  # 增大重叠部分
            length_function=len
        )
        
        splits = text_splitter.split_documents(docs)
        
        if not splits:
            raise ValueError("No text chunks were created from the documents.")

        # Create and store embeddings in Milvus Lite
        try:
            logger.info(f"使用 Milvus Lite 创建向量库，集合名称:{self.collection_name}，数据库文件:{self.milvus_db_file}")
            
            # 使用 MilvusClient 直接连接本地数据库文件
            client = MilvusClient(self.milvus_db_file)
            
            # 检查集合是否存在，如果不存在则创建
            collections = client.list_collections()
            if self.collection_name not in collections:
                # 创建集合
                # bge-small-en 或 m3e-base 模型的维度，这里假设是768
                dim = 768
                client.create_collection(
                    collection_name=self.collection_name,
                    dimension=dim
                )
                logger.info(f"成功创建集合: {self.collection_name}")
            else:
                logger.info(f"集合 {self.collection_name} 已存在，直接添加向量数据")
            
            # 为每个文本块生成嵌入向量并插入到Milvus
            data = []
            for i, doc in enumerate(splits):
                # 生成嵌入向量
                vector = self.embeddings.embed_documents([doc.page_content])[0]
                
                # 构建插入数据
                data_entry = {
                    "id": i,  # 唯一ID
                    "vector": vector,  # 嵌入向量
                    "text": doc.page_content,  # 文本内容
                    "metadata": doc.metadata  # 元数据
                }
                data.append(data_entry)
            
            # 批量插入数据
            if data:
                insert_result = client.insert(
                    collection_name=self.collection_name,
                    data=data
                )
                logger.info(f"成功插入 {len(data)} 条向量数据")
                logger.debug(f"插入结果: {insert_result}")
            
            return "✅ Vector DB Successfully Updated with New Document Data!"
            
        except Exception as e:
            error_msg = f"创建 Milvus Lite 向量存储时发生错误: {e}"
            logger.error(error_msg)
            raise Exception(f"Error creating Milvus vector store: {e}")

    def clear_database(self):
        """
        清空当前向量数据库（删除并重新创建集合）
        
        Returns:
            str: 操作结果消息
        """
        try:
            logger.info(f"正在清空向量数据库，集合名称: {self.collection_name}")
            
            # 连接到数据库
            client = MilvusClient(self.milvus_db_file)
            
            # 检查集合是否存在
            collections = client.list_collections()
            if self.collection_name in collections:
                # 删除集合
                client.drop_collection(self.collection_name)
                logger.info(f"已删除集合: {self.collection_name}")
                
                # 重新创建集合
                dim = 768
                client.create_collection(
                    collection_name=self.collection_name,
                    dimension=dim
                )
                logger.info(f"已重新创建集合: {self.collection_name}")
                
                return "✅ 向量数据库已成功清空！"
            else:
                logger.warning(f"集合 {self.collection_name} 不存在，无需清空。")
                return "⚠️ 向量数据库不存在，无需清空。"
                
        except Exception as e:
            error_msg = f"清空向量数据库时发生错误: {e}"
            logger.error(error_msg)
            return f"⚠️ {error_msg}"

